#include "StdAfx.h"
#include "ExtractData2.h"
#include <iostream>
#include <sstream>
#include <windows.h>
#include <fstream>
#include <direct.h>

//////////////////////////////////////////////////////////////////////////
//Author	:	Ross Conroy ross.conroy@tees.ac.uk
//Date		:	13/11/2014
//
//This is an updated replay extractor that takes in arguments about
//what units to follow and time step duration with the aim of preventing
//the excessively large and impossible to parse output files from old 
//extractor
//////////////////////////////////////////////////////////////////////////

ExtractData2::ExtractData2(void)
{
}

///////////////////////////////////////////////////////////////////////////
//Extracts data about given units only if one of the unit types
//////////////////////////////////////////////////////////////////////////
ExtractData2::ExtractData2(string outputDirectory, vector<string> unitTypes, int timeStepDuration)
{
	_outputDirectory = outputDirectory;
	_unitTypes = unitTypes;
	_timeStepDuration = timeStepDuration;
	_countdown = timeStepDuration;

	stringstream ssDir;
	ssDir << _outputDirectory << "\\";

	for(set<Player*>::iterator it = Broodwar->getPlayers().begin(); it != Broodwar->getPlayers().end(); it++)
	{
		Player * player = *it;
		ssDir << "(" << player->getName() << ")";
	}

	_outputDirectory = ssDir.str();

	cout << "Directory - " << _outputDirectory << endl;
	mkdir(_outputDirectory.c_str());
	_lastTime = clock();
}

//Used at the end of each time step to move the temp rows to their final place for
//later saving
void ExtractData2::CurrentRowsToFinalRows()
{
	for(hash_map<int, DataRow>::iterator it = _currentRows.begin(); it!= _currentRows.end(); it++)
	{
		DataRow row = it->second;
		row.gameTime = _lastTime;
		_finalRows.push_back(row);
	}

	_currentRows.clear();
}

//For each in game tick will get all units and see if any are of interest for extraction
void ExtractData2::Tick()
{
	clock_t currentTime = clock();

	//Countdown time step
	float deltaTime = (currentTime - _lastTime) / ( CLOCKS_PER_SEC / 1000 );
	_countdown = _countdown - deltaTime;
	if (_countdown <= 0)
	{
		_countdown = _timeStepDuration;
		CurrentRowsToFinalRows();
	}

	set<Unit*> units = Broodwar->getAllUnits();

	for(set<Unit*>::const_iterator uit = units.begin(); uit!=units.end(); uit++)
	{
		Unit * unit = *uit;
		for(vector<string>::iterator sit = _unitTypes.begin(); sit != _unitTypes.end(); sit++)
		{
			string unitType = *sit;
			//check if of specified unit types, ignore all others
			if (unit->getType().getName() == unitType)
			{
				Position pos = unit->getPosition();
				int radius = unit->getType().sightRange()*1.5;

				//Green circle = max range, red circle = visible range
				Broodwar->drawCircle(CoordinateType::Map, pos.x(), pos.y(), radius, Colors::Green, false);
				Broodwar->drawCircle(CoordinateType::Map, pos.x(), pos.y(), unit->getType().sightRange(), Colors::Red, false);

				if(_currentRows.find(unit->getID()) == _currentRows.end())
				{
					//Create new row with default values then add to map
					DataRow row;
					row.unitID = unit->getID();
					row.action = NOTHING;
					row.observation = DONT_SEE;
					row.state = NOT_UNDER_ATTACK;
					row.isLandmark = false;
					row.unitHealth = unit->getHitPoints();
					row.unitType = unit->getType().getName();
					_currentRows.insert(pair<int, DataRow>(unit->getID(), row));
				}

				//Determine Action, if not attacking then must be escape or nothing
				if(_currentRows[unit->getID()].action != ATTACK)
				{
					if(unit->isAttackFrame() || unit->isAttacking())
					{
						_currentRows[unit->getID()].action = ATTACK;
					}
				}else if (_currentRows[unit->getID()].action != ESCAPE)
				{
					if((_currentRows[unit->getID()].posX != unit->getPosition().x()) || (_currentRows[unit->getID()].posY != unit->getPosition().y()))
					{
						_currentRows[unit->getID()].action = ESCAPE;
					}
				}

				//Determine state, if lost any health then must have been attacked
				if(_currentRows[unit->getID()].state != UNDER_ATTACK)
				{
					if(unit->isUnderAttack() || unit->getHitPoints() < _currentRows[unit->getID()].unitHealth)
					{
						_currentRows[unit->getID()].state = UNDER_ATTACK;
					}
				}

				//Determine observation, the closest observation overrides the past observation Don't See < See Far < See Near
				if(_currentRows[unit->getID()].observation != SEE_NEAR)
				{
					//Get units within visible radius, if any are enemies get their distance
					float maxDist = -1;
					set<Player*> players = Broodwar->getPlayers();

					for(set<Player*>::iterator itp = players.begin(); itp != players.end(); itp++)
					{
						Player * player = *itp;
						if (player->getID() != unit->getPlayer()->getID())
						{
							set<Unit*> tempUnits = player->getUnits();

							for(set<Unit*>::iterator ittu = tempUnits.begin(); ittu != tempUnits.end(); ittu++)
							{
								Unit * tempUnit = *ittu;
								int tempDist = unit->getDistance(tempUnit);
								if (tempDist <= radius)
								{
									if (maxDist != -1)
									{
										if(tempDist < maxDist)
										{
											maxDist = tempDist;
										}
									}
									else
									{
										maxDist = tempDist;
									}
								}
							}
						}
					}

					if(maxDist != -1)
					{
						if(maxDist < (unit->getType().sightRange()/2))
						{
							_currentRows[unit->getID()].observation = SEE_NEAR;
						}
						else if(maxDist < unit->getType().sightRange())
						{
							_currentRows[unit->getID()].observation = SEE_FAR;
						}
					}

					//use distance to determine battle landmark
					if (maxDist != -1)
					{
						if(maxDist < (unit->getType().sightRange()/2))
						{
							_currentRows[unit->getID()].battleLandmark = CLOSE;
						}
						else if (maxDist < unit->getType().sightRange())
						{
							_currentRows[unit->getID()].battleLandmark = BattleLandMarkDist::MID;
						}
						else
						{
							_currentRows[unit->getID()].battleLandmark = FAR_AWAY;
						}
					}
					else
					{
						_currentRows[unit->getID()].battleLandmark = FAR_AWAY;
					}					
				}				

				if(_currentRows[unit->getID()].action == NOTHING || _currentRows[unit->getID()].action == ESCAPE)
				{
					_currentRows[unit->getID()].status = UnitStatus::IDLE;
					
					//Neutral targets are resources and same player targets are repairs/workers/transporters
					if(unit->getTarget() != NULL)
					{
						if (unit->getTarget()->getPlayer()->isNeutral() || (unit->getTarget()->getPlayer()->getName() == unit->getPlayer()->getName()))
						{
							_currentRows[unit->getID()].status = WORKING;
						}
					}					
					else if((_currentRows[unit->getID()].observation != DONT_SEE) || (_currentRows[unit->getID()].state = UNDER_ATTACK)) //If not working and seen an enemy unit then must be in battle mode
					{
						_currentRows[unit->getID()].status = IN_BATTLE;
					}
				}
				
				if (_currentRows[unit->getID()].action == ATTACK)
				{
					_currentRows[unit->getID()].status = IN_BATTLE;
				}

				if(unit->getHitPoints() < (unit->getType().maxHitPoints() / 3))	
				{
					_currentRows[unit->getID()].landmarkHealth = LOW;
				}
				else if(unit->getHitPoints() < (2 * (unit->getType().maxHitPoints() / 3)))
				{
					_currentRows[unit->getID()].landmarkHealth = LandmarkHealth::MID_HEALTH;
				}
				else
				{
					_currentRows[unit->getID()].landmarkHealth = HIGH;
				}

				_currentRows[unit->getID()].posX = unit->getPosition().x();
				_currentRows[unit->getID()].posY = unit->getPosition().y();
				_currentRows[unit->getID()].unitHealth = unit->getHitPoints();
			}
		}
	}

	_lastTime = currentTime;
}

//Saves all the extracted replay data
void ExtractData2::Save()
{
	CurrentRowsToFinalRows();	

	for(list<DataRow>::iterator it = _finalRows.begin(); it != _finalRows.end(); it++)
	{
		DataRow row = *it;

		if (row.status == IN_BATTLE)
		{
			stringstream ssFile;
			ssFile << _outputDirectory << "\\" << row.unitType << "_" << row.unitID << ".csv";

			bool exists = false;
			ifstream f(ssFile.str().c_str());
			if (f.good()) 
			{
				exists =  true;
			}
			f.close();

			ofstream unitFile;
			unitFile.open (ssFile.str().c_str(), ios_base::app);
			string action;
			string observation;
			string state;
			string status;
			string landmarkDist;
			string landmarkHealth;

			if(row.action == ATTACK)
			{
				action = "Attack";
			}else if(row.action == ESCAPE)
			{
				action = "Escape";
			}else
			{
				action = "Nothing";
			}

			if(row.observation == DONT_SEE)
			{
				observation = "DontSee";
			}else if(row.observation == SEE_FAR)
			{
				observation = "SeeFar";
			}else
			{
				observation = "SeeNear";
			}

			if(row.state == NOT_UNDER_ATTACK)
			{
				state = "NotUnderAttack";
			}else
			{
				state = "UnderAttack";
			}

			if (row.status == IN_BATTLE)
			{
				status = "InBattle";
				if (row.battleLandmark == FAR_AWAY)
				{
					landmarkDist = "Far";
				}
				else if (row.battleLandmark == MID)
				{
					landmarkDist = "Mid";
				}
				else
				{
					landmarkDist = "Close";
				}
			}else if(row.status == WORKING)
			{
				status = "Working";
			}
			else
			{
				status = "Idle";
			}

			switch(row.landmarkHealth)
			{
			case HIGH: landmarkHealth = "High";
				break;
			case MID_HEALTH: landmarkHealth = "Mid";
				break;
			case LOW: landmarkHealth = "Low";
				break;
			default: landmarkHealth = "ERROR";
			}

			if (!exists)
			{
				unitFile << "GameTime,UnitID,UnitType,Action,Observation,State,UnitStatus,LandmarkDist,LandmarkHeath," << endl;
			}

			unitFile << row.gameTime << "," << row.unitID << "," << row.unitType << "," << action << "," << observation << "," << state << "," << status << "," << landmarkDist << "," << landmarkHealth << "," << endl;

			unitFile.close();
		}
	}

	cout << "Save" << endl;
}

